package cn.xxi.ai.application.util;

import cn.xxi.ai.application.dto.ChatContext;
import cn.xxi.ai.application.dto.SseChatMessage;
import cn.xxi.ai.application.entity.Message;
import com.github.yitter.idgen.YitIdHelper;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Sinks;

import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;

/**
 * ChatSinksHolder 用于管理每个会话的消息推送 Sinks。
 * Sinks 是 Project Reactor 中支持发布-订阅推送的组件，适用于 SSE、WebSocket 等。
 * 
 * 支持：
 * - 获取或创建 Sinks
 * - 清理 Sinks
 * - 包装执行逻辑的 use 方法（类似 Kotlin 的 `use {}` 范式）
 * 
 * @author bingo
 */
@Component
public class ChatSinksHolder {

    // 存储 sinksId 与 Sinks.Many 的映射
    private static final ConcurrentHashMap<Long, Sinks.Many<SseChatMessage>> sinksMap = new ConcurrentHashMap<>();

    /**
     * 获取已存在的 Sinks，或如果不存在就创建新的 Sinks。
     *
     * @param sinksId 唯一标识
     * @return Sinks.Many 实例
     */
    public static Sinks.Many<SseChatMessage> getOrCreate(Long sinksId) {
        return sinksMap.computeIfAbsent(sinksId, id ->
                Sinks.many().multicast().onBackpressureBuffer()
        );
    }

    /**
     * 清除指定 ID 的 Sinks，释放资源。
     *
     * @param sinksId 要移除的 sinks ID
     */
    public static void remove(Long sinksId) {
        Sinks.Many<SseChatMessage> sink = sinksMap.get(sinksId);
        if (sink != null) {
            sink.tryEmitComplete();
        }
        sinksMap.remove(sinksId);
    }

    /**
     * 使用 ChatContext 执行一个带返回值的逻辑，自动管理 Sinks 的创建和清除。
     *
     * @param conversationId 会话 ID
     * @param parentId       上级消息 ID，可为空
     * @param block          要执行的逻辑，接收 ChatContext
     * @param <T>            返回值类型
     * @return block 的执行结果
     */
    public static <T> T use(Long conversationId, Long parentId, Function<ChatContext, T> block) {
            long sinksId = YitIdHelper.nextId();
        try {
            getOrCreate(sinksId);
            return block.apply(new ChatContext(conversationId, parentId, sinksId));
        } catch (Exception ex) {
            remove(sinksId);
            throw ex;
        }
    }

    /**
     * 保存并推送消息，持久化部分先用接口占位
     *
     * @param sink       Sinks.Many<SseChatMessage> 实例
     * @param message    消息对象
     * @param forceSave  是否强制保存
     */
    public static void saveAndEmit(Sinks.Many<SseChatMessage> sink, Message message, boolean forceSave) {
        // 生成有ID的新消息，如果当前没有ID
        Message messageWithId = (message.getId() != null) ? message : message.copyWithNewId(YitIdHelper.nextId());

        // TODO: 这里调用持久化接口，暂时注释或实现空方法
        // if (forceSave || messageWithId.getData().persistent()) {
        //     YourPersistenceService.save(messageWithId);
        // }

        // 推送消息到 sink
        sink.tryEmitNext(new SseChatMessage(messageWithId));
    }

    public static void saveAndEmit(Sinks.Many<SseChatMessage> sink, Message message) {
        saveAndEmit(sink, message, false);
    }
}
